#---------------------------------------------------------------#
#    U-Net 
#        ver.2 - ResNet_UNet
#                改进之处 conv -> ResNet @decoder
#---------------------------------------------------------------# 

import tensorflow as tf
import tensorflow.keras as keras

import numpy as np

'''
'''


def BasicBlock(x, channel, strides=(1,1), expansion=1):
    resdual = keras.layers.Conv2D(channel, (3,3), strides=strides, padding='same', use_bias=False)(x)
    resdual = keras.layers.BatchNormalization()(resdual)
    resdual = keras.layers.Activation('relu')(resdual)
    resdual = keras.layers.Conv2D(channel*expansion, (3,3), padding='same', use_bias=False)(resdual)

    shortcut = x
    if strides != (1,1) or x.shape[-1] != channel*expansion:
        shortcut = keras.layers.Conv2D(channel*expansion, (1,1), strides=strides, use_bias=False)(shortcut)
        shortcut = keras.layers.BatchNormalization()(shortcut)

    y = resdual + shortcut
    y = keras.layers.Activation('relu')(y)

    return y

def BottleNeck(x, channel, strides=(1,1), expansion=4):
    residual = keras.layers.Conv2D(channel, (1,1), use_bias=False)(x)
    residual = keras.layers.BatchNormalization()(residual)
    residual = keras.layers.Activation('relu')(residual)
    residual = keras.layers.Conv2D(channel, (3,3), strides=strides, padding='same', use_bias=False)(residual)
    residual = keras.layers.BatchNormalization()(residual)
    residual = keras.layers.Activation('relu')(residual)
    residual = keras.layers.Conv2D(channel*expansion, (1,1), use_bias=False)(residual)
    residual = keras.layers.BatchNormalization()(residual)

    shortcut = x
    if strides != (1,1) or x.shape[-1] != channel*expansion:
        shortcut = keras.layers.Conv2D(channel*expansion, (1,1), strides=strides, use_bias=False)(shortcut)
        shortcut = keras.layers.BatchNormalization()(shortcut)

    y = resdual + shortcut
    y = keras.layers.Activation('relu')(y)

    return y

def ResNet_process(x, block, num_block, num_classes=99):
    '''
        block   ResNet的block过程，为函数
    '''
    assert len(num_block) == 4, 'The given "num_block" is not proper: {}'.format(num_block)

    # feature_map 1
    f1 = keras.layers.Conv2D(64, (7,7), strides=(2,2), padding='same', use_bias=False)(x)
    f1 = keras.layers.BatchNormalization()(f1)
    f1 = keras.layers.Activation('relu')(f1)

    # feature_map 2
    f2 = keras.layers.MaxPooling2D(pool_size=(3,3), strides=(2,2), padding='same')(f1)
    tmp = [1] * (num_block[0] - 1)
    tmp.insert(0, 1)
    for i in tmp:
        f2 = block(f2, channel=64, strides=(i,i))

    # feature_map 3
    tmp = [1] * (num_block[1] - 1)
    tmp.insert(0, 2)
    for i in tmp:
        f3 = block(f2, channel=128, strides=(i,i))

    # feature_map 4
    tmp = [1] * (num_block[2] - 1)
    tmp.insert(0, 2)
    for i in tmp:
        f4 = block(f3, channel=256, strides=(i,i))

    # feature_map 5
    tmp = [1] * (num_block[3] - 1)
    tmp.insert(0, 2)
    for i in tmp:
        f5 = block(f4, channel=512, strides=(i,i))

    y = keras.layers.AveragePooling2D(pool_size=(1,1))(f5)
    y = keras.layers.Flatten()(y)
    y = keras.layers.Dense(num_classes+1)(y)

    return f1, f2, f3, f4, f5, y

# def unet_conv_block(x, out_channels, kernel_size=(3,3), padding='valid', batch_norm_mode=False):
#     y = keras.layers.Conv2D(out_channels, kernel_size, padding=padding, activation='relu', kernel_initializer='he_normal')(x)
#     # if batch_norm_mode:
#     #     y = keras.layers.BatchNormalization()(y) # 注意：BN要在Activation之前做
#     # y = keras.layers.Activation('relu')(y)

#     y = keras.layers.Conv2D(out_channels, kernel_size, padding=padding, activation='relu', kernel_initializer='he_normal')(y)
#     # if batch_norm_mode:
#     #     y = keras.layers.BatchNormalization()(y) # 注意：BN要在Activation之前做
#     # y = keras.layers.Activation('relu')(y)

#     return y

def unet_up_block(x1, x2, out_channels, kernel_size=(2,2), padding='valid', strides=2, up_mode='upconv'):
    '''
        x1 before upconv
        x2 encoder_map
    '''
    assert up_mode in ['upconv', 'upsample'], 'Wrong up_mode input: {}'.format(up_mode)

    if up_mode == 'upconv':
        y = keras.layers.Conv2DTranspose(out_channels, kernel_size, strides, padding)(x1)
    elif up_mode == 'upsample':
        y = keras.layers.UpSampling2D(size=kernel_size, interpolation='bilinear')(x1)
        y = keras.layers.Conv2D(out_channels, (1,1), padding='valid')(y)
    
    _, h1, w1, _ = y.shape
    _, h2, w2, _ = x2.shape
    dh, dw = int((h2-h1)/2), int((w2-w1)/2)
    x2 = x2[:, dh:h1+dh, dw:w1+dw, :]
    assert x2.shape[1:] == y.shape[1:], 'Mismatched shapes of x2 and y in decoder process!'

    y = tf.concat([y, x2], axis=-1)

    # y = unet_conv_block(y, out_channels, (3,3), padding, False)
    #F 改进
    y = BasicBlock(y, out_channels) 
    
    return y

def e_d_process(photo, num_classes): # encoder-decoder process
    assert photo.shape[1:] == [572,572,3], 'Inappropriate input shape: {}'.format(photo.shape)

    channels = [64, 128, 256, 512, 1024]

    ### encoder
    # o = photo
    # # encoder_maps = []
    # for i in range(5):
    #     o = unet_conv_block(o, channels[i])
    #     if i < 4:
    #         encoder_maps.append(o)
    #         o = keras.layers.MaxPooling2D((2,2))(o) 

    # assert o.shape[1:] == [28,28,1024], 'Wrong encoder process! '
    #F 改进
    f1, f2, f3, f4, f5, _ = ResNet_process(photo, block=BasicBlock, num_block=[2,2,2,2], num_classes=num_classes) # resnet_18
    o = f5
    encoder_maps = [f1, f2, f3, f4]

    ### decoder
    for i in range(4):
        k = -1 * (i+1)
        o = unet_up_block(o, encoder_maps[k], channels[k-1], up_mode='upconv')

    return o

def get_net_model(net_input_shape, num_classes):
    assert net_input_shape == [572, 572, 3]
    # 1.输入层
    i_put = keras.layers.Input(net_input_shape)

    # 2.encoder + decoder
    decode_map = e_d_process(i_put, num_classes)

    # 3.输出层
    o_put = keras.layers.Conv2D(num_classes, (1,1), padding='valid')(decode_map)
    o_put = keras.layers.Softmax()(o_put)

    # 4.模型实例
    net_model = keras.models.Model(i_put, o_put)

    return net_model

if __name__ == '__main__':
    # 测试搭建的网络
    net_input_shape = [572, 572, 3]
    num_classes     = 1
    net_model = get_net_model(net_input_shape, num_classes)
    net_model.summary()

    # net_input_shape.insert(0, 2)
    # x = np.random.random(net_input_shape) # 增加batch轴
    # y = net_model(x)
    # y = y.numpy()
    # print('测试网络的前向计算功能：')
    # print('    x.shape = {}'.format(x.shape))
    # print('    y.shape = {}'.format(y.shape))